package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

var (
	httpMethod = map[string]bool{
		"GET":    true,
		"POST":   true,
		"PUT":    true,
		"OPTION": true,
		"DELETE": true,
	}
	httpResp = map[string]bool{
		"HTTP":  true,
		"HTTP2": true,
		"HTTPS": true,
	}
	httpFilter = map[string]bool{}
	handle     *pcap.Handle
	err        error
)

func init() {
	for s, ok := range httpMethod {
		httpFilter[s] = ok
	}
	for s, ok := range httpResp {
		httpFilter[s] = ok
	}
}

var (
	pcapFilePathFlag = flag.String("pcap-path", "", "pcap file path.")
	backServiceFlag  = flag.Bool("b", false, "backstage service.")
	printFlag        = flag.Bool("p", false, "print file text.")
)

func main() {
	flag.Parse()
	if *backServiceFlag {
		http.HandleFunc("/", IndexHandler)
		fmt.Println("启动服务.")
		panic(http.ListenAndServe("0.0.0.0:80", nil))
	}
	if *pcapFilePathFlag == "" {
		log.Println("path empty.")
		return
	}
	filePath := *pcapFilePathFlag

	fp := strings.Split(filepath.Base(filePath), ".")
	if len(fp) != 2 || (len(fp) == 2 && fp[1] != "pcap") {
		log.Println("pcap file error.")
		return
	}

	handle, err = pcap.OpenOffline(filePath)
	if err != nil {
		log.Fatal(err)
	}
	defer handle.Close()

	var filter string = "tcp"
	err = handle.SetBPFFilter(filter)
	if err != nil {
		log.Fatal(err)
	}

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	allBody := &strings.Builder{}
	saveFileName := fp[0] + `_pcap_deal.txt`

	f, err := os.Create(saveFileName)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		f.Close()
		if !*printFlag {
			fmt.Println(saveFileName)
		} else {
			fmt.Println(allBody.String())
		}
	}()

	cache := make(map[string]*strings.Builder)
	for packet := range packetSource.Packets() {
		printPacketInfo(packet, cache)
	}

	for _, builder := range cache {
		writeContent(f, allBody, builder.String(), *printFlag)
	}
}

func printPacketInfo(packet gopacket.Packet, cache map[string]*strings.Builder) {
	ethernetLayer := packet.Layer(layers.LayerTypeEthernet)
	if ethernetLayer != nil {
		ethernetPacket, _ := ethernetLayer.(*layers.Ethernet)
		if ethernetPacket.EthernetType.String() == "IPv4" {
			ipLayer := packet.Layer(layers.LayerTypeIPv4)
			if ipLayer != nil {
				ip, _ := ipLayer.(*layers.IPv4)

				tcpLayer := packet.Layer(layers.LayerTypeTCP)
				if tcpLayer != nil {
					tcp, _ := tcpLayer.(*layers.TCP)
					applicationLayer := packet.ApplicationLayer()

					//fmt.Println("--------------------------------------------------------------------")
					//fmt.Println("Source MAC: ", ethernetPacket.SrcMAC)
					//fmt.Println("Destination MAC: ", ethernetPacket.DstMAC)

					//fmt.Println("Ethernet type: ", ethernetPacket.EthernetType)
					//fmt.Printf("From %s to %s\n", ip.SrcIP, ip.DstIP)
					//fmt.Println("Protocol: ", ip.Protocol)
					//fmt.Printf("From port %d to %d\n", tcp.SrcPort, tcp.DstPort)
					//fmt.Println("Sequence number: ", tcp.Seq)
					key := compound(ip.SrcIP.String(), tcp.SrcPort, ip.DstIP.String(), tcp.DstPort)
					var (
						packetCache *strings.Builder
						ok          bool
					)

					if packetCache, ok = cache[key]; !ok {
						packetCache = &strings.Builder{}
						cache[key] = packetCache
					}

					if applicationLayer != nil {
						payload := string(applicationLayer.Payload())
						for s, _ := range httpFilter {
							if strings.HasPrefix(payload, s) {
								packetCache.WriteString(payload)
								packetCache.WriteString("----------")
								packetCache.WriteByte('\n')
								return
							}
						}
					}
				}
			}
		}
	}
}

func IndexHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintln(w, "hello world")
}

func compound(srcIP string, srcPort layers.TCPPort, dstIP string, dstPort layers.TCPPort) string {
	result := strings.Compare(srcIP, dstIP)
	if result > 0 {
		return fmt.Sprintf("%s%d%s%d", srcIP, srcPort, dstIP, dstPort)
	} else {
		return fmt.Sprintf("%s%d%s%d", dstIP, dstPort, srcIP, srcPort)
	}
}
func writeContent(f *os.File, allBody *strings.Builder, body string, printFlag bool) {
	f.WriteString(fmt.Sprintf("----------------一次http请求总览---------------"))
	f.Write([]byte{'\n'})
	f.WriteString(body)
	f.WriteString("----------------------------------")
	f.Write([]byte{'\n'})
	if printFlag {
		allBody.WriteString(fmt.Sprintf("----------------一次http请求总览---------------"))
		allBody.Write([]byte{'\n'})
		allBody.WriteString(body)
		allBody.WriteString("----------------------------------")
		allBody.Write([]byte{'\n'})
	}
}
